from Crypto.Cipher import AES
import json, traceback
import base64
from . import auth
from flask import request, current_app, jsonify
from flask_login import  login_user
from ..exceptions import ValidationError
from ..models import User
from .. import db
import requests as web_requests
from weixin import WXAPPAPI



@auth.route('/login', methods=['GET'])
def login():
    code = request.args.get('code') if request.args.get('code') else 'null'
    encryptedData = request.args.get('encryptedData') if request.args.get('encryptedData') else 'null'
    iv = request.args.get('iv') if request.args.get('iv') else 'null'

    print('[L] user login')
    print('[*] code: ' + code)

    # 用js_code, appid, secret, grant_type向微信服务器获取session_key, openid, expires_in
    api = WXAPPAPI(appid=current_app.config['APP_ID'],
                  app_secret=current_app.config['SECRET_KEY'])
    session_info = api.exchange_code_for_session_key(code=code)

    # 获取session_info 后
    session_key = session_info.get('session_key')
    openId = session_info.get('openid')

    if session_key and openId:
        print('[*] session_key: %s' % session_key)
        print('[*] openId: %s' % openId)
        # 根据openid判断是否插入用户
        try:
            user = User.query.filter_by(openId=openId).first()
        except Exception as e:
            db.session.rollback()
            print(e)
            response = jsonify(error=traceback.format_exc())
            response.status_code = 403
            return response
        # 如果用户不存在，新建用户
        if user is None:
            # 解密数据
            # encrypted_data 包括敏感数据在内的完整用户信息的加密数据
            # iv 加密算法的初始向量
            decrypted = decrypt(current_app.config['APP_ID'], session_key, encryptedData, iv)

            print('[+] add user: %s' % decrypted)
            user = User(
                openId=decrypted['openId'],
                nickName=decrypted['nickName'],
                gender=decrypted['gender'],
                city=decrypted['city'],
                avatarUrl=decrypted['avatarUrl']
            )

            try:
                db.session.add(user)
                db.session.commit()
            except Exception as e:
                db.session.rollback()
                print(e)
                response = jsonify(error=traceback.format_exc())
                response.status_code = 403
                return response

        # 登录用户，并返回由openid和SECRET_KEY构成的token // 3rd_session
        login_user(user, True)
        response = jsonify(message='Login successfully!', user_info=user.to_json())
        response.status_code = 200
        return response

    return str(session_info), 400


# 解密加密数据，获取watermark中的appid
def decrypt(appid, session_key, encryptedData, iv):
    sessionKey = base64.b64decode(session_key)
    encryptedData = base64.b64decode(encryptedData)
    iv = base64.b64decode(iv)

    cipher = AES.new(sessionKey, AES.MODE_CBC, iv)
    s = cipher.decrypt(encryptedData)
    s = str(s, encoding='utf-8')
    decrypted = json.loads(s[:-ord(s[len(s)-1:])])
    print(decrypted)
    if decrypted['watermark']['appid'] != appid:
        raise Exception('Invalid Buffer')

    return decrypted